package handler

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"net/http"
	"strings"
	"time"
	"github.com/golang-jwt/jwt/v4"

	"project-info/src/model"
	"project-info/src/service"

	"project-info/src/lib")

// RSA私钥用于JWT签名（实际应该从配置文件读取）
var rsaPrivateKey *rsa.PrivateKey
var rsaPublicKey *rsa.PublicKey

// 初始化RSA密钥对
func init() {
	// 初始化OIDC服务
	oidcService = service.NewOIDCService()

	// 从配置文件读取RSA私钥
	initRSAKeys()
}

// initRSAKeys 初始化RSA密钥
func initRSAKeys() {
	oidcConfig := model.GloOIDC

	if oidcConfig.SigningKey != "" {
		// 从配置文件解析RSA私钥
		block, _ := pem.Decode([]byte(oidcConfig.SigningKey))
		if block == nil {
			panic("Failed to parse PEM block containing the private key")
		}

		privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			// 尝试PKCS8格式
			parsedKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
			if err != nil {
				panic("Failed to parse RSA private key: " + err.Error())
			}
			var ok bool
			privateKey, ok = parsedKey.(*rsa.PrivateKey)
			if !ok {
				panic("Not an RSA private key")
			}
		}

		rsaPrivateKey = privateKey
		rsaPublicKey = &privateKey.PublicKey
	} else {
		// 如果配置文件中没有私钥，生成临时密钥对
		privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			panic("Failed to generate RSA key pair: " + err.Error())
		}
		rsaPrivateKey = privateKey
		rsaPublicKey = &privateKey.PublicKey
	}
}

// GetRSAPublicKey 获取RSA公钥（供其他包使用）
func GetRSAPublicKey() *rsa.PublicKey {
	return rsaPublicKey
}

// Token 令牌端点
// POST /oidc/token
func Token(c *lib.GinContext) (err error) {
	// 获取grant_type
	grantType := c.PostForm("grant_type")
	if grantType == "" {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_request",
			"error_description": "Missing grant_type parameter",
		})
	}

	switch grantType {
	case "authorization_code":
		return handleAuthorizationCodeGrant(c)
	case "refresh_token":
		return handleRefreshTokenGrant(c)
	default:
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "unsupported_grant_type",
			"error_description": "Unsupported grant type: " + grantType,
		})
	}
}

// handleAuthorizationCodeGrant 处理授权码授权
func handleAuthorizationCodeGrant(c *lib.GinContext) error {
	// 获取请求参数
	code := c.PostForm("code")
	clientID := c.PostForm("client_id")
	clientSecret := c.PostForm("client_secret")
	redirectURI := c.PostForm("redirect_uri")
	codeVerifier := c.PostForm("code_verifier")

	// 验证必需参数
	if code == "" {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_request",
			"error_description": "Missing code parameter",
		})
	}
	if clientID == "" {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_request",
			"error_description": "Missing client_id parameter",
		})
	}

	// 验证客户端
	var client *model.OIDCClient
	var err error
	if clientSecret != "" {
		// 使用客户端密钥验证
		client, err = oidcService.ValidateClient(clientID, clientSecret)
	} else {
		// 仅验证客户端ID（用于公共客户端）
		client, err = oidcService.GetClient(clientID)
	}
	if err != nil {
		return c.JSON(http.StatusUnauthorized, map[string]interface{}{
			"error":             "invalid_client",
			"error_description": err.Error(),
		})
	}

	// 验证授权码
	authCodeData, err := oidcService.GetAuthorizationCode(code)
	if err != nil {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_grant",
			"error_description": err.Error(),
		})
	}

	// 验证客户端ID匹配
	if authCodeData.ClientID != clientID {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_grant",
			"error_description": "Client ID mismatch",
		})
	}

	// 验证重定向URI匹配
	if redirectURI != "" && authCodeData.RedirectURI != redirectURI {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_grant",
			"error_description": "Redirect URI mismatch",
		})
	}

	// 验证PKCE（如果使用）
	if authCodeData.CodeChallenge != "" {
		if codeVerifier == "" {
			return c.JSON(http.StatusBadRequest, map[string]interface{}{
				"error":             "invalid_request",
				"error_description": "Missing code_verifier parameter",
			})
		}

		codeChallengeMethod := authCodeData.CodeChallengeMethod
		if codeChallengeMethod == "" {
			codeChallengeMethod = "plain"
		}

		if !verifyPKCE(codeVerifier, authCodeData.CodeChallenge, codeChallengeMethod) {
			return c.JSON(http.StatusBadRequest, map[string]interface{}{
				"error":             "invalid_grant",
				"error_description": "Invalid code_verifier",
			})
		}
	}

	// 标记授权码为已使用
	err = oidcService.UseAuthorizationCode(code)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to mark authorization code as used",
		})
	}

	// 生成访问令牌
	accessToken, err := generateAccessToken()
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to generate access token",
		})
	}

	// 生成刷新令牌
	refreshToken, err := generateRefreshToken()
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to generate refresh token",
		})
	}

	// 生成ID令牌
	idToken, err := generateIDToken(client, authCodeData)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to generate ID token",
		})
	}

	// 创建访问令牌记录
	now := time.Now()
	accessTokenRecord := &model.OIDCAccessToken{
		Token:     accessToken,
		ClientID:  clientID,
		UserID:    authCodeData.UserID,
		UserName:  authCodeData.UserName,
		Scope:     authCodeData.Scope,
		ExpiresAt: now.Add(time.Duration(client.AccessTokenTTL) * time.Second),
	}

	// 保存访问令牌到数据库
	err = oidcService.CreateAccessToken(accessTokenRecord)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to store access token",
		})
	}

	// 创建刷新令牌记录
	refreshTokenRecord := &model.OIDCRefreshToken{
		Token:         refreshToken,
		AccessTokenID: accessTokenRecord.Id,
		ClientID:      clientID,
		UserID:        authCodeData.UserID,
		UserName:      authCodeData.UserName,
		Scope:         authCodeData.Scope,
		ExpiresAt:     now.Add(time.Duration(client.RefreshTokenTTL) * time.Second),
	}

	// 保存刷新令牌到数据库
	err = oidcService.CreateRefreshToken(refreshTokenRecord)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to store refresh token",
		})
	}

	// 返回令牌响应
	response := map[string]interface{}{
		"access_token":  accessToken,
		"token_type":    "Bearer",
		"expires_in":    client.AccessTokenTTL,
		"refresh_token": refreshToken,
		"scope":         authCodeData.Scope,
	}

	// 如果请求包含openid scope，添加ID令牌
	if strings.Contains(authCodeData.Scope, "openid") {
		response["id_token"] = idToken
	}

	return c.JSON(http.StatusOK, response)
}

// handleRefreshTokenGrant 处理刷新令牌授权
func handleRefreshTokenGrant(c *lib.GinContext) error {
	// 获取请求参数
	refreshToken := c.PostForm("refresh_token")
	clientID := c.PostForm("client_id")
	clientSecret := c.PostForm("client_secret")
	scope := c.PostForm("scope")

	// 验证必需参数
	if refreshToken == "" {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_request",
			"error_description": "Missing refresh_token parameter",
		})
	}
	if clientID == "" {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_request",
			"error_description": "Missing client_id parameter",
		})
	}

	// 验证客户端
	var client *model.OIDCClient
	var err error
	if clientSecret != "" {
		// 使用客户端密钥验证
		client, err = oidcService.ValidateClient(clientID, clientSecret)
	} else {
		// 仅验证客户端ID（用于公共客户端）
		client, err = oidcService.GetClient(clientID)
	}
	if err != nil {
		return c.JSON(http.StatusUnauthorized, map[string]interface{}{
			"error":             "invalid_client",
			"error_description": err.Error(),
		})
	}

	// 验证刷新令牌
	refreshTokenData, err := oidcService.GetRefreshToken(refreshToken)
	if err != nil {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_grant",
			"error_description": err.Error(),
		})
	}

	// 验证客户端ID匹配
	if refreshTokenData.ClientID != clientID {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_grant",
			"error_description": "Client ID mismatch",
		})
	}

	// 验证作用域（如果提供）
	if scope != "" && !isScopeSubset(scope, refreshTokenData.Scope) {
		return c.JSON(http.StatusBadRequest, map[string]interface{}{
			"error":             "invalid_scope",
			"error_description": "Requested scope exceeds original scope",
		})
	}

	// 使用原始作用域（如果没有提供新的作用域）
	if scope == "" {
		scope = refreshTokenData.Scope
	}

	// 生成新的访问令牌
	newAccessToken, err := generateAccessToken()
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to generate access token",
		})
	}

	// 创建新的访问令牌记录
	now := time.Now()
	newAccessTokenRecord := &model.OIDCAccessToken{
		Token:     newAccessToken,
		ClientID:  clientID,
		UserID:    refreshTokenData.UserID,
		UserName:  refreshTokenData.UserName,
		Scope:     scope,
		ExpiresAt: now.Add(time.Duration(client.AccessTokenTTL) * time.Second),
	}

	// 保存新的访问令牌到数据库
	err = oidcService.CreateAccessToken(newAccessTokenRecord)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]interface{}{
			"error":             "server_error",
			"error_description": "Failed to store access token",
		})
	}

	// 返回令牌响应
	response := map[string]interface{}{
		"access_token": newAccessToken,
		"token_type":   "Bearer",
		"expires_in":   client.AccessTokenTTL,
		"scope":        scope,
	}

	// 如果原始刷新令牌仍然有效，可以继续使用
	// 在某些实现中，可能会生成新的刷新令牌
	response["refresh_token"] = refreshToken

	return c.JSON(http.StatusOK, response)
}

// generateAccessToken 生成访问令牌
func generateAccessToken() (string, error) {
	bytes := make([]byte, 32)
	_, err := rand.Read(bytes)
	if err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(bytes), nil
}

// generateRefreshToken 生成刷新令牌
func generateRefreshToken() (string, error) {
	bytes := make([]byte, 32)
	_, err := rand.Read(bytes)
	if err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(bytes), nil
}

// generateIDToken 生成ID令牌
func generateIDToken(client *model.OIDCClient, authCodeData *model.OIDCAuthorizationCode) (string, error) {
	now := time.Now()
	oidcConfig := model.GloOIDC

	// 创建JWT声明
	claims := jwt.MapClaims{
		"iss":       oidcConfig.Issuer,                                              // 发行者
		"sub":       authCodeData.UserID,                                            // 主题（用户ID）
		"aud":       client.ClientID,                                                // 受众（客户端ID）
		"exp":       now.Add(time.Duration(client.IDTokenTTL) * time.Second).Unix(), // 过期时间
		"iat":       now.Unix(),                                                     // 签发时间
		"auth_time": now.Unix(),                                                     // 认证时间
	}

	// 添加nonce（如果存在）
	if authCodeData.Nonce != "" {
		claims["nonce"] = authCodeData.Nonce
	}

	// 根据作用域添加用户信息
	scopes := strings.Fields(authCodeData.Scope)
	scopeConfigs, err := oidcService.GetScopes(scopes)
	if err != nil {
		return "", err
	}

	var user model.User
	err = model.DB().Where("id = ? and deleted = 0", authCodeData.UserID).First(&user).Error
	if err != nil {
		return "", err
	}

	// 2. 动态构建声明
	for _, scopeConfig := range scopeConfigs {
		for _, claim := range scopeConfig.Claims {
			value := getClaimValue(user, claim)
			if value != nil {
				claims[claim] = value
			}
		}
	}

	// 创建JWT令牌
	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)

	// 设置Key ID（从配置文件读取）
	token.Header["kid"] = oidcConfig.KeyId

	// 签名令牌
	tokenString, err := token.SignedString(rsaPrivateKey)
	if err != nil {
		return "", err
	}

	return tokenString, nil
}

// verifyPKCE 验证PKCE
func verifyPKCE(codeVerifier, codeChallenge, codeChallengeMethod string) bool {
	switch codeChallengeMethod {
	case "plain":
		return codeVerifier == codeChallenge
	case "S256":
		hash := sha256.Sum256([]byte(codeVerifier))
		expected := base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString(hash[:])
		return expected == codeChallenge
	default:
		return false
	}
}

// isScopeSubset 检查请求的作用域是否是原始作用域的子集
func isScopeSubset(requestedScope, originalScope string) bool {
	requestedScopes := strings.Fields(requestedScope)
	originalScopes := strings.Fields(originalScope)

	originalScopeMap := make(map[string]bool)
	for _, scope := range originalScopes {
		originalScopeMap[scope] = true
	}

	for _, scope := range requestedScopes {
		if !originalScopeMap[scope] {
			return false
		}
	}

	return true
}
